# -*- coding: UTF-8 -*-
import json
from typing import Dict, Optional

import aioredis


class YRedis(aioredis.Redis):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    async def hset_dict(self, key: str, field: str = None, value: str = None, field_value: Dict = None,
                        expire: int = None, nx: bool = False):
        """
        通过字典传入创建redis哈希类型，同时设置缓存过期时间（原生的模块不支持，使用lua脚本实现）
        :param key: 主键
        :param field: 子健
        :param value: 子值
        :param field_value: 字典类型的键值 {子健: 子值}
        :param expire: 过期时间
        :param nx: 是否只在键不存在时，才对键（主键）进行设置操作
        :return:
        """
        key_dict = {
            "key": key,
            "expire": "NOT EXPIRE",
            "nx": "NOT NX",
            "xx": "NOT XX"
        }
        if int(expire) > 0:
            key_dict["expire"] = int(expire)

        if nx:
            key_dict["nx"] = "NX"

        fields = []
        if field and value:
            fields = [field, value]

        for k, v in field_value.items():
            fields.append(k)
            fields.append(json.dumps(v, ensure_ascii=False, separators=(',', ':')))

        # redis 是单线程不用加锁
        lua_str = """
            local key = KEYS[1]
            local expire = KEYS[2]
            local nx = KEYS[3]
            local value = ARGV
            
            if nx == "NX" then cmd = "hsetnx" else cmd = "hset" end
            
            for i=1, #value-1, 2 do
                redis.call(cmd, key, value[i], value[i+1])
            end    
            
            if expire == "NOT EXPIRE" then 
                redis.call('expire', key)
            else 
                redis.call('expire', key, expire)
            end
            """

        script = self.register_script(lua_str)
        await script(keys=list(key_dict.values()), args=fields)

    async def yhscan(self, name: str, cursor: int = 0, match: Optional[str] = "*", count: Optional[int] = None):
        # todo 待测试 scan 和 scan_iter https://zhuanlan.zhihu.com/p/150790074
        # todo 参数校验是不是也应该自己封装一个通用的？
        """
        获取全部 hash 的键值对
        :param name: 主键
        :param cursor: 从当前游标开始查起
        :param match: 匹配规则
        :param count:
        :return:
        """

        field_value = {}
        count, data = await self.hscan(name, cursor, match, count)  # todo Redis是不是已经做过处理了，直接获取的就是全部？这个count好像也没什么用
        field_value.update(data)

        while count:
            count, data = await self.hscan(name, cursor, match, count)
            field_value.update(data)

        return field_value

    async def yscan(self, match: Optional[str] = "*", limit: int = None, cursor: int = 0):
        """
        扫描全表获取指定数量的set键值
        :param match: redis匹配规则
        :param limit: 返回数据数量在 limit 之内 这里有坑， scan 有可能重复获取，count参数不好使，有可能一次获取全部为空，但是游标不是0，所以暂时先获取所有键，在limit取值
        :param cursor: 游标
        :return: # fixme 键、值是中文时会显示乱码，
        """
        fields = set()
        while True:
            cursor, data = await self.scan(cursor, match)
            fields.update(set(data))

            if not cursor:
                break

        fields = list(fields) if limit and len(fields) <= limit else list(fields)[:limit]
        return fields


async def get_redis_connection():
    redis_conf = {
        "username": "",
        "password": "",
        "host": "localhost",
        "port": "6379",
        "db": "0",
        "max_connections": 8
    }
    return YRedis(**redis_conf)

